import os
import uuid
import operator
from typing import List, Annotated, Literal
from langchain_community.chat_models import ChatTongyi
from langchain_core.messages import SystemMessage, AIMessage, HumanMessage, ToolMessage
from pydantic import BaseModel
from langgraph.graph import StateGraph, END, START
from langgraph.checkpoint.memory import MemorySaver


"""
好使
"""
# 工具数据模型
class SearchToolInput(BaseModel):
    query: str

# 工具函数
def search_tool(args: SearchToolInput):
    # 这里可以接入真实搜索API

    return f"搜索结果：{args.query} 的百科内容……"

# 初始化 LLM
llm = ChatTongyi(model="qwen-max", api_key=os.getenv("aliyun_API_KEY"), temperature=0)
llm_with_tool = llm.bind_tools([SearchToolInput])

def get_messages_info(messages):
    return [SystemMessage(content="你是一个智能助手，可以调用搜索工具。用户有问题时请主动调用工具。")] + messages

# 对话链
def chain(messages):
    return [llm_with_tool.invoke(get_messages_info(messages))]

# 工具调用节点
def tool_node(messages):
    ai_msg = messages[-1]
    tool_args = ai_msg.tool_calls[0]["args"]
    result = search_tool(SearchToolInput(**tool_args))
    return [ToolMessage(content=result, tool_call_id=ai_msg.tool_calls[0]["id"])]

# 状态转移
def get_state(messages) -> Literal["tool", "chat", "__end__"]:
    if isinstance(messages[-1], AIMessage) and messages[-1].tool_calls:
        return "tool"
    elif not isinstance(messages[-1], HumanMessage):
        return END
    return "chat"

# 初始化 StateGraph
memory = MemorySaver()
workflow = StateGraph(Annotated[list, operator.add])
workflow.add_node("chat", chain)
workflow.add_node("tool", tool_node)
workflow.add_conditional_edges("chat", get_state)
workflow.add_edge("tool", "chat")
workflow.add_edge(START, "chat")
workflow.add_edge("chat", END)
graph = workflow.compile(checkpointer=memory)

# 交互主循环
config = {"configurable": {"thread_id": str(uuid.uuid4())}}
while True:
    user = input("User (q/Q to quit): ")
    if user in {"q", "Q"}:
        print("AI: Byebye")
        break
    output = None
    for output in graph.stream([HumanMessage(content=user)], config=config, stream_mode="updates"):
        last_message = next(iter(output.values()))
        print(last_message)
    if output and "chat" in output:
        print("Done!")